import json

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

from starlette.responses import StreamingResponse

from RAGFlow_output import create_session, stream_question_generator

app = FastAPI()

# 配置CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class QuestionRequest(BaseModel):
    question: str

@app.post("/LLM/RAGFlow")
async def stream_api(request: QuestionRequest) -> StreamingResponse:
    """FastAPI流式接口"""
    question = request.question.strip()
    if not question:
        raise HTTPException(status_code=400, detail="问题参数不能为空")

    session_id = create_session(question)
    if not session_id:
        raise HTTPException(status_code=500, detail="会话创建失败")

    # 流式响应生成器
    def generate():

        for fragment in stream_question_generator(question, session_id):

            if "</think>" in fragment:
                # 拆分片段为:前、</think>标记、后三部分
                parts = fragment.split("</think>", 1)  # 最多分割一次
                before = parts[0]
                after = parts[1] if len(parts) > 1 else ""
                # 返回前半部分（如果非空）
                if before.strip():
                    yield f"data: {json.dumps({'type': 'RAGFlow', 'content': before})}\n\n"
                # 返回</think>标记
                yield f"data: {json.dumps({'type': 'RAGFlow', 'content': '</think>'})}\n\n"
                # 返回后半部分（如果非空）
                if after.strip():
                    yield f"data: {json.dumps({'type': 'RAGFlow', 'content': after})}\n\n"
            else:
                yield f"data: {json.dumps({'type': 'RAGFlow', 'content': fragment})}\n\n"

            print(fragment, end='', flush=True)

    return StreamingResponse(generate(), media_type="text/event-stream")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=3001)